from langchain_core.messages import HumanMessage
from typing import Literal
from langgraph.prebuilt import ToolNode
from langgraph.graph import END, StateGraph, MessagesState
import toolsLibrary as tL
import llmShift
from langgraph.checkpoint.mysql.pymysql import PyMySQLSaver
import sqlConnecter

tools = [tL.search_images, tL.get_transit_route, tL.get_walking_route, tL.get_scenic_spots, tL.get_douyin_hot]

tool_node = ToolNode(tools)

model = llmShift.version_choose().bind_tools(tools)


def should_continue(state: MessagesState) -> Literal["tools", END]:
    messages = state['messages']
    last_message = messages[-1]
    if last_message.tool_calls:
        return 'tools'
    return END


def call_model(state: MessagesState):
    messages = state['messages']
    response = model.invoke(messages)
    return {"messages": [response]}


workflow = StateGraph(MessagesState)
workflow.set_entry_point("agent")  # 入口节点

workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)
# workflow.add_node("check_response", check_response)

workflow.add_conditional_edges("agent", should_continue, )
workflow.add_edge("tools", 'agent')

checkpointer = PyMySQLSaver(sqlConnecter.connecter)
checkpointer.setup()

app = workflow.compile(checkpointer=checkpointer)


# def AGENT(query, thread_id=11):
#     final_state = app.invoke(
#         {"messages": [HumanMessage(content=f"{query}")]},
#         config={"configurable": {"thread_id": thread_id}}
#     )
#     result = final_state["messages"][-1].content
#     return result

def AGENT(query, thread_id=9):
    for chunk in app.stream(
            {"messages": [HumanMessage(content=f"{query}")]},
            stream_mode="values",
            config={"configurable": {"thread_id": thread_id}}
    ):
        chunk["messages"][-1].pretty_print()


print(AGENT('今天的抖音热点视频话题有哪些？'))
